import datasets
import json
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv
import re
from logger import get_logger

load_dotenv()

llm = ChatOpenAI(temperature=0.1, model="/file/ljw22/Qwen2.5-7B-Instruct")

muser = datasets.load_from_disk("MUSER-main/muser-dataset")

logger = get_logger()
# muser = muser.select(range(1))

with open("./explain.json") as fin:
    explain_dict = json.load(fin)


def make_prompt(fact, atom, explain, event):
    prompt = """\
下面给你一段放在三引号中的关于“{event_type}”的`法律文书事实描述`：
```
{fact}
```
请你根据上述的事实描述，提取出其中的“{atom}”。
“{atom}”指的是：{explain}。
分析并提取出来之后，再以“【@你的答案@】”的格式输出，以便于我进行提取。如果不存在，就在你的答案部分输出“不存在“，格式还是要保持。
"""

    template = ChatPromptTemplate([
        ("system", "你是一个精通法律知识的助手，可以帮助人们分析法律文书，并能准确的提取出所需要的信息。"),
        ("human", prompt),
    ])
    
    prompt_value = template.invoke(
        {
            "event_type": event,
            "fact": fact,
            "atom": atom,
            "explain": explain
        }
    )
    return prompt_value


def extract_from_resp(text):
    # print("----" * 10)
    # print(text)
    # print("----" * 10)
    # 定义正则表达式
    pattern = r"【@([^@]+)@】"

    # 使用正则表达式搜索文本
    match = re.search(pattern, text)

    # 如果找到匹配项，则提取并输出
    if match:
        answer = match.group(1)  # 提取匹配的内容
        return answer
    else:
        return "不存在"

fout = open("atom_res_dict.json", '+a')

def process(data):
    event_type_list = data["event_type"]
    fact = data["fact"]
    
    result_dict = {}
    atom_dict = {}
    idx = data["idx"]
    index = data["index"]
    
    result_dict["index"] = index
    result_dict["idx"] = idx
    
    for event in event_type_list:
        atom_dict[event] = {}
        """
        抽取出 fact 事实中的，事件类型为 event 下的所有的要素原子
        """
        res = "不存在"
        for atom, explain in explain_dict[event].items():
            prompt = make_prompt(fact, atom, explain, event)
            try:
                resp = llm.invoke(prompt).content
                res = extract_from_resp(resp)
            except Exception as e:
                logger.info(f"第【{index}】个数据，事件【{event}】，原子【{atom}】抽取失败，数据编号【{index}】")
                logger.info("异常：")
                logger.info(str(e))
            atom_dict[event][atom] = res
    result_dict["atom"] = atom_dict
    fout.write(json.dumps(result_dict, ensure_ascii=False) + '\n')
    fout.flush()

annotated_muser = muser.map(process)
fout.close()